package com.demo.spark.sql;

import org.apache.spark.SparkConf;
import org.apache.spark.api.java.JavaRDD;
import org.apache.spark.api.java.JavaSparkContext;
import org.apache.spark.api.java.function.ForeachFunction;
import org.apache.spark.api.java.function.Function;
import org.apache.spark.api.java.function.MapFunction;
import org.apache.spark.sql.*;
import org.apache.spark.sql.catalyst.expressions.Encode;
import org.apache.spark.sql.types.DataTypes;
import org.apache.spark.sql.types.StructField;
import org.apache.spark.sql.types.StructType;

import java.util.ArrayList;
import java.util.List;

public class DataFrameSchema {
    public static void main(String[] args) {
        SparkConf sparkConf = new SparkConf().setMaster("local").setAppName("data frame ");
        JavaSparkContext sparkContext = new JavaSparkContext(sparkConf);
        sparkContext.setLogLevel("error");
        SQLContext sqlContext = new SQLContext(sparkContext);

        String schemaString = "name age";

        // Generate the schema based on the string of schema
        List<StructField> fields = new ArrayList<>();
        for (String fieldName : schemaString.split(" ")) {
            StructField field = DataTypes.createStructField(fieldName, DataTypes.StringType, true);
            fields.add(field);
        }
        StructType schema = DataTypes.createStructType(fields);
        JavaRDD<Row> map = sparkContext.textFile("D:\\git\\hadoop-cloudera-demo\\src\\main\\resources\\person.txt").map(new Function<String, Row>() {
            @Override
            public Row call(String v1) throws Exception {

                String[] s = v1.split(" ");
                Row row = RowFactory.create(s[0], s[1]);
                return row;
            }
        });
        Dataset<Row> dataFrame = sqlContext.createDataFrame(map, schema);
        dataFrame.createOrReplaceTempView("person");
        Dataset<Row> resultdata = sqlContext.sql("select * from person");
        Dataset<String> map1 = resultdata.map(new MapFunction<Row, String>() {
            @Override
            public String call(Row value) throws Exception {
                return value.get(0) + "," + value.get(1);
            }
        }, Encoders.STRING());
        map1.foreach(new ForeachFunction<String>() {
            @Override
            public void call(String s) throws Exception {
                System.out.println(s);
            }
        });

        //resultdata.write().format("avro").save("./test/person.avro");
        resultdata.write().format("parquet").save("./test/person.parquet");
        //resultdata.write().format("json").save("./test/person.json");
        //resultdata.write().format("csv").save("./test/person.csv");
        //resultdata.write().format("orc").save("./test/person.orc");
    }
}
